package ml.classification;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;

import ml.WekaUtil;
import ml.classification.filter.PostProcessingFilter;
import weka.classifiers.Classifier;
import weka.core.Attribute;
import weka.core.Instances;

public class Classification {

	private static final String NEW_LINE = "\r\n";

	private static ArrayList<Integer> listTrueClassMusic;
	private static ArrayList<Integer> listPredictedMusic;
	private static ArrayList<Integer> listPostprocessedMusic;

	private static ArrayList<Integer> listTrueClassSpeech;
	private static ArrayList<Integer> listPredictedSpeech;
	private static ArrayList<Integer> listPostprocessedSpeech;

	private static ArrayList<Integer> listTrueClassCombined;
	private static ArrayList<Integer> listPredictedCombined;
	private static ArrayList<Integer> listPostprocessedCombined;
	private static ArrayList<Integer> listSpeechAlwaysTrueCombined;

	private static ArrayList<Double> listMaximumGain;
	private static ArrayList<Double> listPredictedGain;
	private static ArrayList<Double> listPostprocessedGain;
	private static ArrayList<Double> listSpeechAlwaysTrueGain;

	private static Classifier musicClassifier;
	private static Classifier speechClassifier;

	private static Instances musicTestData;
	private static Instances speechTestData;

	private static int numInstances;

	public static ClassificationStatistics classifyAndPostprocess(
			final String fileName, final boolean outputFiles,
			final String outputDirectory, final boolean isContestData,
			final String musicTestDataFileName,
			final String speechTestDataFileName,
			final String musicModelFileName, final String speechModelFileName,
			final PostProcessingFilter musicFilter,
			final PostProcessingFilter speechFilter) throws Exception {

		final long startMs = System.currentTimeMillis();

		loadClassifiers(musicModelFileName, speechModelFileName);

		loadMusicAndSpeechData(musicTestDataFileName, speechTestDataFileName);

		initLists();

		classifyMusic(isContestData);

		classifySpeech(isContestData);

		classifyCombinedClassAndCalculateGain();

		postProcessMusic(musicFilter);

		postProcessSpeech(speechFilter);

		classifyCombinedClassAndCalculateGainForPostProcessedData(isContestData);

		final ClassificationStatistics stats = calculateStatistics(fileName,
				musicFilter, speechFilter, musicModelFileName,
				speechModelFileName);

		if (outputFiles) {
			writeFileForVisualization(fileName, outputDirectory,
					musicTestDataFileName, speechTestDataFileName,
					musicModelFileName, speechModelFileName);

			writeMusicTextFileForSubmission(fileName, outputDirectory);

			writeSpeechTextFileForSubmission(fileName, outputDirectory);
		}

		System.out.println("\n");
		System.out.println(" - Maximum Gain:         " + stats.maximumGain);
		System.out.println(" - Gain via prediction:  " + stats.predictedGain);
		System.out.println(" - Gain after postproc.: "
				+ stats.postprocessedGain);
		System.out.println(" - Gain Speech Always 1: "
				+ stats.speechAlwaysTrueGain);
		System.out.println("----------------------------------------");
		System.out.println(" - musicN: " + musicFilter.getN());
		System.out.println(" - musicIterations: "
				+ musicFilter.getNoIterations());
		System.out.println(" - speechN: " + speechFilter.getN());
		System.out.println(" - speechIterations: "
				+ speechFilter.getNoIterations());
		System.out.println("----------------------------------------");
		System.out.println(" --> Finished classification for '" + fileName
				+ "'   (" + (System.currentTimeMillis() - startMs) + " ms).");

		deinitLists();

		return stats;
	}

	private static ClassificationStatistics calculateStatistics(
			final String fileName, final PostProcessingFilter musicFilter,
			final PostProcessingFilter speechFilter,
			final String musicModelFileName, final String speechModelFileName) {
		// statistics
		final ClassificationStatistics stats = new ClassificationStatistics();

		stats.fileName = fileName;

		// sum up gain
		for (int i = 0; i < numInstances; i++) {
			stats.maximumGain += listMaximumGain.get(i);
			stats.predictedGain += listPredictedGain.get(i);
			stats.postprocessedGain += listPostprocessedGain.get(i);
			stats.speechAlwaysTrueGain += listSpeechAlwaysTrueGain.get(i);
		}

		// count number of classes
		int predictedCorrectClass = 0;
		int postprocessedCorrectClass = 0;
		int predictedCorrectMusic = 0;
		int postprocessedCorrectMusic = 0;
		int predictedCorrectSpeech = 0;
		int postprocessedCorrectSpeech = 0;

		for (int i = 0; i < numInstances; i++) {
			// true class
			if (listTrueClassCombined.get(i) == 0) {
				stats.trueClassNoMusic++;
			} else if (listTrueClassCombined.get(i) == 1) {
				stats.trueClassMusicSpeech++;
			} else if (listTrueClassCombined.get(i) == 2) {
				stats.trueClassMusicOnly++;
			}
			if (listTrueClassMusic.get(i) == 1) {
				stats.trueClassMusic++;
			}
			if (listTrueClassSpeech.get(i) == 1) {
				stats.trueClassSpeech++;
			}
			// predicted class
			if (listPredictedCombined.get(i) == 0) {
				stats.predictedClassNoMusic++;
			} else if (listPredictedCombined.get(i) == 1) {
				stats.predictedClassMusicSpeech++;
			} else if (listPredictedCombined.get(i) == 2) {
				stats.predictedClassMusicOnly++;
			}
			if (listPredictedMusic.get(i) == 1) {
				stats.predictedClassMusic++;
			}
			if (listPredictedSpeech.get(i) == 1) {
				stats.predictedClassSpeech++;
			}
			// postprocessed class
			if (listPostprocessedCombined.get(i) == 0) {
				stats.postprocessedClassNoMusic++;
			} else if (listPostprocessedCombined.get(i) == 1) {
				stats.postprocessedClassMusicSpeech++;
			} else if (listPostprocessedCombined.get(i) == 2) {
				stats.postprocessedClassMusicOnly++;
			}
			if (listPostprocessedMusic.get(i) == 1) {
				stats.postprocessedClassMusic++;
			}
			if (listPostprocessedSpeech.get(i) == 1) {
				stats.postprocessedClassSpeech++;
			}
			// count correct classified instances
			if (listTrueClassCombined.get(i) == listPredictedCombined.get(i)) {
				predictedCorrectClass++;
			}
			if (listTrueClassCombined.get(i) == listPostprocessedCombined
					.get(i)) {
				postprocessedCorrectClass++;
			}
			if (listTrueClassMusic.get(i) == listPredictedMusic.get(i)) {
				predictedCorrectMusic++;
			}
			if (listTrueClassMusic.get(i) == listPostprocessedMusic.get(i)) {
				postprocessedCorrectMusic++;
			}
			if (listTrueClassSpeech.get(i) == listPredictedSpeech.get(i)) {
				predictedCorrectSpeech++;
			}
			if (listTrueClassSpeech.get(i) == listPostprocessedSpeech.get(i)) {
				postprocessedCorrectSpeech++;
			}
		}
		// compute accuracy
		stats.predictedAccuracyClass = predictedCorrectClass
				/ (double) numInstances;
		stats.postprocessedAccuracyClass = postprocessedCorrectClass
				/ (double) numInstances;
		stats.predictedAccuracyMusic = predictedCorrectMusic
				/ (double) numInstances;
		stats.postprocessedAccuracyMusic = postprocessedCorrectMusic
				/ (double) numInstances;
		stats.predictedAccuracySpeech = predictedCorrectSpeech
				/ (double) numInstances;
		stats.postprocessedAccuracySpeech = postprocessedCorrectSpeech
				/ (double) numInstances;

		// append configuration parameters
		stats.musicFilter = musicFilter.getFilterName();
		stats.musicN = musicFilter.getN();
		stats.musicIterations = musicFilter.getNoIterations();
		stats.speechFilter = speechFilter.getFilterName();
		stats.speechN = speechFilter.getN();
		stats.speechIterations = speechFilter.getNoIterations();
		stats.speechThreshold = speechFilter.getThreshold();

		// data file
		stats.musicModelFileName = musicModelFileName;
		stats.speechModelFileName = speechModelFileName;

		return stats;
	}

	private static void writeSpeechTextFileForSubmission(final String fileName,
			final String outputDirectory) throws IOException {
		// write text file for submission including speech classification
		System.out.print("\n - Write speech classification data: ");
		final BufferedWriter writerSpeech = new BufferedWriter(new FileWriter(
				outputDirectory + fileName + "_speech.txt"));
		for (int i = 0; i < numInstances; i++) {
			writerSpeech.write(String.valueOf(listPostprocessedSpeech.get(i)));
			writerSpeech.write(NEW_LINE);
			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
		writerSpeech.flush();
		writerSpeech.close();
	}

	private static void writeMusicTextFileForSubmission(final String fileName,
			final String outputDirectory) throws IOException {
		// write text file for submission including music classification
		System.out.print("\n - Write music classification data:  ");
		final BufferedWriter writerMusic = new BufferedWriter(new FileWriter(
				outputDirectory + fileName + "_music.txt"));

		for (int i = 0; i < numInstances; i++) {
			writerMusic.write(String.valueOf(listPostprocessedMusic.get(i)));
			writerMusic.write(NEW_LINE);
			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
		writerMusic.flush();
		writerMusic.close();
	}

	private static void writeFileForVisualization(final String fileName,
			final String outputDirectory, final String musicTestDataFileName,
			final String speechTestDataFileName,
			final String musicModelFileName, final String speechModelFileName)
			throws IOException {
		// write data to output file
		System.out.print("\n - Write data: ");
		String fileNameTmp = "speech_music";
		fileNameTmp = fileName.toString(); // TODO
		final BufferedWriter writer = new BufferedWriter(new FileWriter(
				outputDirectory + fileNameTmp + ".csv"));
		final int lineSize = 120 * 5; // 2 Minuten

		writer.write("\"Music test data: ");
		writer.write(musicTestDataFileName);
		writer.write("\"");
		writer.write(NEW_LINE);
		writer.write("\"Speech test data: ");
		writer.write(speechTestDataFileName);
		writer.write("\"");
		writer.write(NEW_LINE);
		writer.write("\"Music classifier model: ");
		writer.write(musicModelFileName);
		writer.write("\"");
		writer.write(NEW_LINE);
		writer.write("\"Speech classifier model: ");
		writer.write(speechModelFileName);
		writer.write("\"");
		writer.write(NEW_LINE);
		writer.write(NEW_LINE);
		for (int part = 0; part < listTrueClassMusic.size() / lineSize; part++) {
			// music class
			writer.write("\"Music True Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listTrueClassMusic.size(); i++) {
				writer.write(String.valueOf(listTrueClassMusic.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Music Predicted Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPredictedMusic.size(); i++) {
				writer.write(String.valueOf(listPredictedMusic.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Music Postprocessed Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPostprocessedMusic.size(); i++) {
				writer.write(String.valueOf(listPostprocessedMusic.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			// speech class
			writer.write("\"Speech True Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listTrueClassSpeech.size(); i++) {
				writer.write(String.valueOf(listTrueClassSpeech.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Speech Predicted Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPredictedSpeech.size(); i++) {
				writer.write(String.valueOf(listPredictedSpeech.get(i)));
				writer.write(";");
			}
			writer.write("\n");
			writer.write("\"Speech Postprocessed Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPostprocessedSpeech.size(); i++) {
				writer.write(String.valueOf(listPostprocessedSpeech.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			// combined class
			writer.write("\"True Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listTrueClassCombined.size(); i++) {
				writer.write(String.valueOf(listTrueClassCombined.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Predicted Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPredictedCombined.size(); i++) {
				writer.write(String.valueOf(listPredictedCombined.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Postprocessed Class\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPostprocessedCombined.size(); i++) {
				writer.write(String.valueOf(listPostprocessedCombined.get(i)));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			// gain
			writer.write("\"Maximum Gain\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listMaximumGain.size(); i++) {
				writer.write(String.valueOf(listMaximumGain.get(i)).replace(
						".", ","));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Predicted Gain\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPredictedGain.size(); i++) {
				writer.write(String.valueOf(listPredictedGain.get(i)).replace(
						".", ","));
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write("\"Postprocessed Gain\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listPostprocessedGain.size(); i++) {
				writer.write(String.valueOf(listPostprocessedGain.get(i))
						.replace(".", ","));
				writer.write(";");
			}
			// time in seconds
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write("\"Time in seconds\";");
			for (int i = part * lineSize; i < (part + 1) * lineSize
					&& i < listTrueClassMusic.size(); i++) {
				if (i % 25 == 0) {
					writer.write(String.valueOf(i / 5));
				}
				writer.write(";");
			}
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			writer.write(NEW_LINE);
			System.out.print(".");
		}
		writer.flush();
		writer.close();
	}

	private static void classifyCombinedClassAndCalculateGainForPostProcessedData(
			final boolean isContestData) {
		// combined classification && gain (for postprocessed data)
		System.out
				.print("\n - Combined classification and gain calculation (for postprocessed data): ");
		for (int i = 0; i < numInstances; i++) {
			final int trueClass = listTrueClassCombined.get(i);
			final int postprocessedClass = combineClassifications(
					listPostprocessedMusic.get(i),
					listPostprocessedSpeech.get(i));
			listPostprocessedCombined.add(postprocessedClass);

			final double postprocessedGain = gain(trueClass, postprocessedClass);
			listPostprocessedGain.add(postprocessedGain);

			if (isContestData) {
				listMaximumGain.set(i,
						gain(postprocessedClass, postprocessedClass));
			}

			// hypothetical: speech always 1
			final int speechAlwaysTrueClass = combineClassifications(
					listPostprocessedMusic.get(i), 1);
			listSpeechAlwaysTrueCombined.add(speechAlwaysTrueClass);
			final double speechAlwaysTrueGain = gain(trueClass,
					speechAlwaysTrueClass);
			listSpeechAlwaysTrueGain.add(speechAlwaysTrueGain);
			if (isContestData) {
				listSpeechAlwaysTrueGain.set(i,
						gain(speechAlwaysTrueClass, speechAlwaysTrueClass));
			}
			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
	}

	private static void postProcessSpeech(
			final PostProcessingFilter speechFilter) {
		// postprocess speech
		System.out.print("\n - Postprocess speech: ");
		listPostprocessedSpeech = speechFilter.apply(listPredictedSpeech);
	}

	private static void postProcessMusic(final PostProcessingFilter musicFilter) {
		// postprocess music
		System.out.print("\n - Postprocess music: ");
		listPostprocessedMusic = musicFilter.apply(listPredictedMusic);
	}

	private static void classifyCombinedClassAndCalculateGain() {
		// combined classification && gain
		System.out.print("\n - Combined classification and gain calculation: ");
		for (int i = 0; i < numInstances; i++) {
			// combined classification
			final int trueClass = combineClassifications(
					listTrueClassMusic.get(i), listTrueClassSpeech.get(i));
			final int predictedClass = combineClassifications(
					listPredictedMusic.get(i), listPredictedSpeech.get(i));
			listTrueClassCombined.add(trueClass);
			listPredictedCombined.add(predictedClass);

			// gain
			final double maximumGain = gain(trueClass, trueClass);
			final double predictedGain = gain(trueClass, predictedClass);
			listMaximumGain.add(maximumGain);
			listPredictedGain.add(predictedGain);

			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
	}

	private static void classifySpeech(final boolean isContestData)
			throws Exception {
		// classify speech
		final long startMs = System.currentTimeMillis();
		System.out.print("\n - Classify speech: ");
		final Attribute speechClassAttr = speechTestData
				.attribute(speechTestData.numAttributes() - 1);
		for (int i = 0; i < numInstances; i++) {
			final int trueClassSpeech = isContestData ? 0
					: ((int) speechTestData.instance(i).value(speechClassAttr));
			final int predictedSpeech = (int) speechClassifier
					.classifyInstance(speechTestData.instance(i));
			listTrueClassSpeech.add(trueClassSpeech);
			listPredictedSpeech.add(predictedSpeech);
			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
		System.out
				.print(" (" + (System.currentTimeMillis() - startMs) + " ms)");
	}

	private static void classifyMusic(final boolean isContestData)
			throws Exception {
		// classify music
		final long startMs = System.currentTimeMillis();
		System.out.print("\n - Classify music:  ");
		final Attribute musicClassAttr = musicTestData.attribute(musicTestData
				.numAttributes() - 1);
		for (int i = 0; i < numInstances; i++) {
			final int trueClassMusic = isContestData ? 1 : ((int) musicTestData
					.instance(i).value(musicClassAttr));
			final int predictedMusic = (int) musicClassifier
					.classifyInstance(musicTestData.instance(i));
			// invert music, no_music
			listTrueClassMusic.add((trueClassMusic == 0) ? 1 : 0);
			// invert music, no_music
			listPredictedMusic.add((predictedMusic == 0) ? 1 : 0);
			if (i % 1000 == 0) {
				System.out.print(".");
			}
		}
		System.out
				.print(" (" + (System.currentTimeMillis() - startMs) + " ms)");
	}

	private static void loadClassifiers(final String musicModelFileName,
			final String speechModelFileName) throws Exception {
		musicClassifier = (Classifier) (weka.core.SerializationHelper
				.read(musicModelFileName));
		speechClassifier = (Classifier) (weka.core.SerializationHelper
				.read(speechModelFileName));
	}

	private static void loadMusicAndSpeechData(
			final String musicTestDataFileName,
			final String speechTestDataFileName) throws Exception {
		musicTestData = WekaUtil.loadData(musicTestDataFileName);
		speechTestData = WekaUtil.loadData(speechTestDataFileName);
		if (musicTestData.numInstances() != speechTestData.numInstances()) {
			throw new Exception(
					"Number of music instances and speech instances does not ");
		}
		numInstances = musicTestData.numInstances();
	}

	private static void initLists() {
		listTrueClassMusic = new ArrayList<Integer>();
		listPredictedMusic = new ArrayList<Integer>();
		listPostprocessedMusic = new ArrayList<Integer>();

		listTrueClassSpeech = new ArrayList<Integer>();
		listPredictedSpeech = new ArrayList<Integer>();
		listPostprocessedSpeech = new ArrayList<Integer>();

		listTrueClassCombined = new ArrayList<Integer>();
		listPredictedCombined = new ArrayList<Integer>();
		listPostprocessedCombined = new ArrayList<Integer>();
		listSpeechAlwaysTrueCombined = new ArrayList<Integer>();

		listMaximumGain = new ArrayList<Double>();
		listPredictedGain = new ArrayList<Double>();
		listPostprocessedGain = new ArrayList<Double>();
		listSpeechAlwaysTrueGain = new ArrayList<Double>();
	}

	private static void deinitLists() {
		listTrueClassMusic = null;
		listPredictedMusic = null;
		listPostprocessedMusic = null;

		listTrueClassSpeech = null;
		listPredictedSpeech = null;
		listPostprocessedSpeech = null;

		listTrueClassCombined = null;
		listPredictedCombined = null;
		listPostprocessedCombined = null;
		listSpeechAlwaysTrueCombined = null;

		listMaximumGain = null;
		listPredictedGain = null;
		listPostprocessedGain = null;
		listSpeechAlwaysTrueGain = null;

		musicTestData = null;
		speechTestData = null;
	}

	private static int combineClassifications(final int classMusic,
			final int classSpeech) {
		if (classMusic < 0 || classMusic > 1) {
			throw new IllegalArgumentException("classMusic must be 0 or 1");
		}
		if (classSpeech < 0 || classSpeech > 1) {
			throw new IllegalArgumentException("classSpeech must be 0 or 1");
		}

		if (classMusic == 1 && classSpeech == 0) {
			// music only
			return 2;
		} else if (classMusic == 1 && classSpeech == 1) {
			// music & speech
			return 1;
		}
		// no music
		return 0;
	}

	private static double gain(final int trueClassMusic,
			final int predictedClassMusic, final int trueClassSpeech,
			final int predictedClassSpeech) throws IllegalArgumentException {
		if (trueClassMusic < 0 || trueClassMusic > 1) {
			throw new IllegalArgumentException("trueClassMusic must be 0 or 1");
		}
		if (predictedClassMusic < 0 || predictedClassMusic > 1) {
			throw new IllegalArgumentException(
					"predictedClassMusic must be 0 or 1");
		}
		if (trueClassSpeech < 0 || trueClassSpeech > 1) {
			throw new IllegalArgumentException("trueClassSpeech must be 0 or 1");
		}
		if (predictedClassSpeech < 0 || predictedClassSpeech > 1) {
			throw new IllegalArgumentException(
					"predictedClassSpeech must be 0 or 1");
		}

		final int trueClass = combineClassifications(trueClassMusic,
				trueClassSpeech);
		final int predictedClass = combineClassifications(predictedClassMusic,
				predictedClassSpeech);

		return gain(trueClass, predictedClass);
	}

	private static double gain(final int trueClass, final int predictedClass)
			throws IllegalArgumentException {
		if (trueClass < 0 || trueClass > 2) {
			throw new IllegalArgumentException("trueClass must be 0, 1 or 2");
		}
		if (predictedClass < 0 || predictedClass > 2) {
			throw new IllegalArgumentException(
					"predictedClass must be 0, 1 or 2");
		}

		if (trueClass == 2) {
			// true class: music only
			if (predictedClass == 2) {
				// predicted class: music only
				return 1.00;
			}
			if (predictedClass == 1) {
				// predicted class: music & speech
				return 0.75;
			}
		} else if (trueClass == 1) {
			// true class: music & speech
			if (predictedClass == 2) {
				// predicted class: music only
				return -0.50;
			}
			if (predictedClass == 1) {
				// predicted class: music & speech
				return 0.75;
			}
		} else if (trueClass == 0) {
			// true class: no music
			if (predictedClass == 2) {
				// predicted class: music only
				return -3.00;
			}
			if (predictedClass == 1) {
				// predicted class: music & speech
				return -1.00;
			}
		}

		// predicted class: no music
		return 0;
	}
}
